'''
第一个RAG修改版本，使用大模型帮忙阅读在线HTML阅读理解
'''

# 导入关键包
import bs4
from langchain import hub
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import WebBaseLoader
from langchain_community.vectorstores import Chroma
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough
from langchain_community.llms import Ollama
from langchain_community.embeddings import OllamaEmbeddings
from langchain.prompts import ChatPromptTemplate

# 加载在线文档
url = "https://mp.weixin.qq.com/s/hk9TfzmzhFuXkuDugr7OqA"
loader = WebBaseLoader(web_path=(url,),)
docs = loader.aload()

# 进行分词
text_splitter = RecursiveCharacterTextSplitter(chunk_size=300, chunk_overlap=50)
splits = text_splitter.split_documents(docs)

# Embed(终端运行此命令下载模型：ollama pull mofanke/acge_text_embedding)
vectorstore = Chroma.from_documents(documents=splits, embedding=OllamaEmbeddings(model="mofanke/acge_text_embedding"))

retriever = vectorstore.as_retriever()
# 完成索引构建
print("index done")

# 使用自定义prompt
# template = """你是一个回答问题的助手。请使用以下检索到的背景信息来回答问题。如果你不知道答案，直接说你不知道。请用最多三句话来保持回答的简洁性。
# 问题：{question}
#
# 背景：{context}
#
# 答案：
# """
template = """你是一个总结文章的助手。请使用以下检索到的背景信息来总结文章。如果你无法总结，直接说你无法总结。请用归纳方法提取核心思想的方法来保持回答的简洁性。"""

# prompt = ChatPromptTemplate.from_template(template=template, input_variables=["context", "question"])
prompt = ChatPromptTemplate.from_template(template)
# 打印出提示词
print(prompt)

llm = Ollama(model="qwen:0.5b")

def format_docs(docs):
    return "\n\n".join(doc.page_content for doc in docs)

# 编排服务的Chain
rag_chain = (
    {"context": retriever | format_docs, "question": RunnablePassthrough()}
    | prompt
    | llm
    | StrOutputParser()
)
# 提出问题
re = rag_chain.get_prompts(rag_chain)
print(re)
result = rag_chain.invoke("习近平赴山东日照考察了什么?")

print(result)


